#include "ASTVisitor.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Verifier.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/GlobalVariable.h"
#include <iostream>

namespace decaf_lang::parser {

std::any ASTVisitor::visitProgram(DecafLangParser::ProgramContext *ctx) {
    // Create main function: int main()
    llvm::FunctionType* funcType = llvm::FunctionType::get(builder.getInt32Ty(), false);
    llvm::Function* mainFunc = llvm::Function::Create(
        funcType, llvm::Function::ExternalLinkage, "main", module);

    llvm::BasicBlock* entryBlock = llvm::BasicBlock::Create(context, "entry", mainFunc);
    builder.SetInsertPoint(entryBlock);

    // Visit children (decl and printStmt)
    for (auto child : ctx->children) {
        visit(child);
    }

    // Return 0
    builder.CreateRet(llvm::ConstantInt::get(context, llvm::APInt(32, 0)));

    // Verify function
    if (llvm::verifyFunction(*mainFunc, &llvm::errs())) {
        llvm::errs() << "Error verifying main function\n";
    }

    return nullptr;
}

std::any ASTVisitor::visitDecl(DecafLangParser::DeclContext *ctx) {
    if (ctx->INT()) {
        std::string varName = ctx->ID()->getText();
        int value = std::stoi(ctx->INT()->getText());
        llvm::AllocaInst* alloca = builder.CreateAlloca(builder.getInt32Ty(), nullptr, varName);
        builder.CreateStore(
            llvm::ConstantInt::get(context, llvm::APInt(32, value)),
            alloca
        );
        intVars[varName] = alloca;
    } else if (ctx->STRING()) {
        std::string varName = ctx->ID()->getText();
        std::string value = ctx->STRING()->getText();
        value = value.substr(1, value.length() - 2); // Remove quotes
        llvm::Value* strPtr = builder.CreateGlobalStringPtr(value);
        stringVars[varName] = strPtr;
    }
    return nullptr;
}

std::any ASTVisitor::visitPrintStmt(DecafLangParser::PrintStmtContext *ctx) {
    auto* expr = ctx->expr();
    llvm::Function* printfFunc = module.getFunction("printf");
    if (!printfFunc) {
        std::vector<llvm::Type*> printfArgs{builder.getInt8PtrTy()};
        llvm::FunctionType* printfType = llvm::FunctionType::get(
            builder.getInt32Ty(), printfArgs, true);
        printfFunc = llvm::Function::Create(
            printfType, llvm::Function::ExternalLinkage, "printf", module);
    }

    if (expr->INT()) {
        int value = std::stoi(expr->INT()->getText());
        llvm::Value* formatStr = builder.CreateGlobalStringPtr("%d\n");
        builder.CreateCall(printfFunc, {formatStr, llvm::ConstantInt::get(context, llvm::APInt(32, value))});
    } else if (expr->STRING()) {
        std::string str = expr->STRING()->getText();
        str = str.substr(1, str.length() - 2);
        llvm::Value* formatStr = builder.CreateGlobalStringPtr(str + "\n");
        builder.CreateCall(printfFunc, {formatStr});
    } else if (expr->BOOL()) {
        std::string str = expr->BOOL()->getText();
        llvm::Value* formatStr = builder.CreateGlobalStringPtr(str + "\n");
        builder.CreateCall(printfFunc, {formatStr});
    } else if (expr->ID()) {
        std::string varName = expr->ID()->getText();
        // Check intVars first
        auto itInt = intVars.find(varName);
        if (itInt != intVars.end()) {
            llvm::Value* varPtr = itInt->second;
            llvm::Value* loaded = builder.CreateLoad(builder.getInt32Ty(), varPtr, varName);
            llvm::Value* formatStr = builder.CreateGlobalStringPtr("%d\n");
            builder.CreateCall(printfFunc, {formatStr, loaded});
        } else {
            // Check stringVars
            auto itStr = stringVars.find(varName);
            if (itStr != stringVars.end()) {
                llvm::Value* formatStr = builder.CreateGlobalStringPtr("%s\n");
                builder.CreateCall(printfFunc, {formatStr, itStr->second});
            } else {
                std::cerr << "Undefined variable: " << varName << std::endl;
            }
        }
    }
    return nullptr;
}

std::any ASTVisitor::visitAssign(DecafLangParser::AssignContext *ctx) {
    std::string varName = ctx->ID()->getText();
    auto* expr = ctx->expr();

    // Handle int assignment
    if (expr->INT()) {
        int value = std::stoi(expr->INT()->getText());
        auto it = intVars.find(varName);
        if (it == intVars.end()) {
            std::cerr << "Undefined int variable: " << varName << std::endl;
            return nullptr;
        }
        llvm::Value* varPtr = it->second;
        builder.CreateStore(
            llvm::ConstantInt::get(context, llvm::APInt(32, value)),
            varPtr
        );
    }
    // Handle string assignment
    else if (expr->STRING()) {
        std::string value = expr->STRING()->getText();
        value = value.substr(1, value.length() - 2); // Remove quotes
        auto it = stringVars.find(varName);
        if (it == stringVars.end()) {
            std::cerr << "Undefined string variable: " << varName << std::endl;
            return nullptr;
        }
        // For simplicity, just update the pointer in the symbol table
        llvm::Value* strPtr = builder.CreateGlobalStringPtr(value);
        stringVars[varName] = strPtr;
    }
    // Handle variable-to-variable assignment (optional, for int/string)
    else if (expr->ID()) {
        std::string rhs = expr->ID()->getText();
        // Try int
        auto itL = intVars.find(varName);
        auto itR = intVars.find(rhs);
        if (itL != intVars.end() && itR != intVars.end()) {
            llvm::Value* rhsPtr = itR->second;
            llvm::Value* loaded = builder.CreateLoad(builder.getInt32Ty(), rhsPtr, rhs);
            builder.CreateStore(loaded, itL->second);
        } else {
            // Try string
            auto itLS = stringVars.find(varName);
            auto itRS = stringVars.find(rhs);
            if (itLS != stringVars.end() && itRS != stringVars.end()) {
                stringVars[varName] = itRS->second;
            } else {
                std::cerr << "Type mismatch or undefined variable in assignment: " << varName << " = " << rhs << std::endl;
            }
        }
    }
    return nullptr;
}


} // namespace decaf_lang::parser
